import torch
import torchvision.datasets
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


class Ah(nn.Module):
    def __init__(self):
        super().__init__()
        self.model1 = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64)
            ,Linear(64,10)
        )

    def forward(self,x):
        x = self.model1(x)
        return x


ah = Ah()
print(ah)


dataSet = torchvision.datasets.CIFAR10("../../dataSet",train=True,
                                       download=True,
                                       transform=torchvision.transforms.ToTensor())

dataLoader = DataLoader(dataSet,64,drop_last=True)


writer = SummaryWriter("../logs")

step = 0
loss = nn.CrossEntropyLoss()
# 1.
lr = 0.001
optim = torch.optim.SGD(ah.parameters(),lr,)
sheduler = StepLR(optim,step_size=5,gamma=0.1)
for epoch in range(100):
    running_loss = 0.0
    for data in dataLoader:
        imgs , tables = data
        output = ah(imgs)
        lossRes = loss(output,tables)
        # she zhi ti du
        optim.zero_grad()
        # huo de te du
        lossRes.backward()
        optim.step()
        #sheduler.step()
        running_loss +=lossRes

    print(running_loss)

writer.close()



